import time

from time import sleep

from queue import Queue

from pymongo import MongoClient

from conf.settings import CLOUDMONITOR_MONGODB, CLOUDMONITOR_DB
from utility.log import Logger
from utility.monitor import Monitor
from utility.utils import RescanTaskThread

logger = Logger().get_logger()


db = MongoClient(CLOUDMONITOR_MONGODB)[CLOUDMONITOR_DB]
# db = MongoClient(host='localhost')['cloudmonitor']

monitor = Monitor()

vul_type_list = [
    "file_upload",
    "directory_listing",
    "xss",
    "directory_traversal",
    "code_execution",
    "csrf",
    "file_inclusion",
    "injection",
    "configuration",
    "information_disclosure",
    "command_execution",
    "malware",
    "webvul"
]

# task_type_tuple = ('vul', 'securityEvent')
# task_type_tuple = ('http', 'ping')
# task_type_tuple = ('securityEvent', 'content', 'vul', 'ssl')
task_type_tuple = ('asset', )


class TaskConfig(object):
    area_map = {
        "chengdu-cmc1": "xiangyang-cmc1",
        "kunming-ctc1": "xian2-ctc1",
        "chongqing-cuc1": "zhuzhou-cuc1",
        "beijing-cuc1": "nanchang2-cuc1",
        "dalian-ctc1": "dalian-ctc1",
        "hengyang-ctc1": "hengyang-ctc1",
    }

    # area_map = {v: k for k, v in area_map.items()}

    def __init__(self, config, task_type):
        self.config = config

        # prefix
        if task_type == 'securityEvent':
            task_type = 'security_event'

        self.func = getattr(self, f"get_{task_type}")

    def get(self):
        return self.func()

    def get_vul(self):
        return self.config

    def get_http(self):
        for item in self.config:
            item.setdefault('collect', {})['area'] = self._replace_area(item.get('collect', {}).get('area', []))
        return self.config

    def get_ping(self):
        return self.get_http()

    def get_security_event(self):
        item = self.config.setdefault("collect", {}).setdefault("plugins", [])
        item.append('foreign_links')
        self.config['collect']['plugins'] = list(set(item))

        return self.config

    def get_content(self):
        return self.config

    def _replace_area(self, old_area):
        new_area = []
        new_area_list = list(self.area_map.values())
        for area in old_area:
            new_area.append(self.area_map.get(area, area))

        new_area = list(set(new_area))
        i = 0
        while len(new_area) < len(old_area):
            new_area.append(new_area_list[i])
            new_area = list(set(new_area))
            i += 1
            if i >= len(new_area_list):
                break
        return new_area


def task_type_rescan(job_query, ignore_check=True):
    size = 10
    q = Queue(size)

    i = 0

    jobs = db['jobs'].find(job_query)
    asset_ids = jobs.distinct('assetId')

    task_type = "asset"
    for asset in db['assets'].find({"_id": {"$in": asset_ids}}):
        config = asset.get("taskSettings")
        if config:
            q.put(str(asset['_id']))

            config['enable'] = True
            config['interval'] = 604800000

            RescanTaskThread(
                db['jobs'], monitor, config, task_type, asset['jobId'], logger, q, ignore_check=ignore_check).start()
            i += 1
            if i % 10 == 0:
                sleep(1)

    q.join()


def main() -> int:
    job_query = {'verification.isVerified': True}

    logger.info(f"Start {__file__}")
    task_type_rescan(job_query)

    logger.info(f"End {__file__}")

    return 0


if __name__ == '__main__':
    main()
